/**
 * Copyright (c) 2009 Koder Chen. All Rights Reserved.
 *
 * rsa.c
 *
 * The file implements the function declared in the rsa.h file
 *
 *  Author: chenfuhan
 * Created: 2009-6-23
 * Version: 1.0
 */
#include <math.h>
#include <time.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <memory.h>
#include <assert.h>

#include "tool.h"
#include "rsa.h"
#include "bigint.h"

u8_pt jacobi_a   = NULL;
u8_pt jacobi_b   = NULL;
u8_pt jacobi_r   = NULL;
u8_pt jacobi_o   = NULL;
u8_pt jacobi_t1  = NULL;
u8_pt jacobi_t2  = NULL;
u8_pt jacobi_x   = NULL;
u8_pt jacobi_min = NULL;
u8_pt jacobi_nt  = NULL;
u8_pt jacobi_nh  = NULL;

void rsa_conv_int2bytes(i32_t a, u8_pt out)
{
	i32_t i = 0;
	while(a > 0)
	{
		out[i] = (a & 0xff);
		a = (a >> 8);
		++i;
	}
}

i32_t rsa_conv_bytes2int(u8_pt in)
{
	return (in[3] << 24) + (in[2] << 16) + (in[1] << 8) + in[0];
}

void rsa_rand_bits(i32_t bits, u8_pt out)
{
	i32_t bytes = ((bits+7) >> 3), i=0;
	assert(bytes <= MAX_BYTES / 2);
	for(i=0;i<bytes;i++)
	{
		i32_t t = rsa_rand();
		out[i] = (t & 0xff);
	}
}

void rsa_rand_int(u8_pt min, u8_pt max, u8_pt out)
{
	i32_t minBits = 8, bits = 0;
	u8_pt range = rsa_default_malloc();

	bi_sub(max, min, range);
	bits = rsa_max(bi_bits(range), minBits);
	// generate the rand number with the specified bits
	u8_pt r = rsa_default_malloc();
	rsa_rand_bits(bits, r);
	bi_mod(r, range, r);
	bi_add(r, min, out);
	//
	rsa_free(range);
	rsa_free(r);
}

i32_t rsa_fermat_test(u8_pt a)
{
	u8_pt r = rsa_default_malloc();
	u8_pt t = rsa_default_malloc();
	i32_t v = 0;

	bi_set_bit(r, 0, 1);
	bi_sub(a, r, t);
	rsa_rand_int(r, t, r);
	bi_powm(r, t, a, r);

	if(bi_bits(r) == 1) v = 1;
	rsa_free(r);
	rsa_free(t);
	return v;
}

i32_t rsa_jacobi(u8_pt m, u8_pt n)
{
	bi_reset(jacobi_o);
	bi_set_bit(jacobi_o, 0, 1);
	u8_pt t = 0;

	bi_mod(m, n, jacobi_r);
	if(bi_bits(jacobi_r) == 0)
	{
		return 0;
	}
	bi_copy_bytes(jacobi_a, m, 0, MAX_BYTES);
	bi_copy_bytes(jacobi_b, n, 0, MAX_BYTES);
	// now the loop
	i32_t bits = bi_bits(jacobi_a), result = 1;
	while(bits > 1)
	{
		if(bi_get_bit(jacobi_a, 0))
		{
			bi_sub(jacobi_a, jacobi_o, jacobi_t1);
			bi_sub(jacobi_b, jacobi_o, jacobi_t2);
			bi_mul(jacobi_t1, jacobi_t2, jacobi_t1);
			bi_rshift(jacobi_t1, 2);
			if(bi_get_bit(jacobi_t1, 0)) result = -result;
			t = jacobi_a;
			bi_mod(jacobi_b, jacobi_a, jacobi_r);
			jacobi_a = jacobi_r;
			jacobi_r = jacobi_b;
			jacobi_b = t;
			t = 0;
		}
		else
		{
			bi_mul(jacobi_b, jacobi_b, jacobi_t1);
			bi_sub(jacobi_t1, jacobi_o, jacobi_t1);
			bi_rshift(jacobi_t1, 3);
			if(bi_get_bit(jacobi_t1, 0)) result = -result;
			bi_rshift(jacobi_a, 1);
		}
		bits = bi_bits(jacobi_a);
	}
	return result;
}

i32_t rsa_jacobi_test(u8_pt n, i32_t k)
{
	i32_t i=0, j=0, retv = 1;
	bi_reset(jacobi_min);
	bi_set_bit(jacobi_min, 0, 1);
	bi_sub(n, jacobi_min, jacobi_nh);

	for(i=0;i<=k;i++)
	{
		bi_copy_bytes(jacobi_nt, jacobi_nh, 0, MAX_BYTES);
		rsa_rand_int(jacobi_min, jacobi_nh, jacobi_x);
		j = rsa_jacobi(jacobi_x, n);
		bi_rshift(jacobi_nt, 1);
		bi_powm(jacobi_x, jacobi_nt, n, jacobi_nt);
		if(j >-1 && bi_bits(jacobi_nt) != j)
		{
			retv = 0;
			break;
		}
		if(j == -1 && (bi_comp(jacobi_nt, jacobi_nh) != 0))
		{
			retv = 0;
			break;
		}
	}
	return retv;
}

i32_t rsa_check_prime(u8_pt a)
{
	return rsa_fermat_test(a);
}

i32_t rsa_check_prime_strong(u8_pt a)
{
	//return (rsa_fermat_test(a) && rsa_jacobi_test(a, 5));
	return rsa_jacobi_test(a, 5);
}

void rsa_rand_prime(i32_t bits, u8_pt out)
{
	while(1)
	{
		rsa_rand_bits(bits, out);
		bi_set_bit(out, bits-1, 1);
		bi_set_bit(out, 0, 1);
		if(rsa_check_prime_strong(out) == 1)
		{
			break;
		}
	}
}

i32_t rsa_relatively_prime(u8_pt a, u8_pt b)
{
	u8_pt t = rsa_default_malloc();
	i32_t retv = 0;
	bi_gcd(a, b, t);
	if(bi_bits(t) == 1)
	{
		retv = 1;
	}
	rsa_free(t);
	return retv;
}

void rsa_gen_p_q(i32_t bits, u8_pt p, u8_pt q)
{
	rsa_rand_prime(bits, p);
	while(1)
	{
		rsa_rand_prime(bits, q);
		if(rsa_relatively_prime(p, q) == 1) return;		
	}
}

void rsa_extended_euclid_gcd(u8_pt a, u8_pt b, u8_pt d, u8_pt k, u8_pt l, i32_pt markX, i32_pt markY)
{
	if(bi_bits(b) == 0)
	{
		bi_reset(k);
		bi_reset(l);
		bi_set_bit(k, 0, 1);
		bi_copy_bytes(d, a, 0, MAX_BYTES);
		*markX = 0;
		*markY = 0;
		return;
	}
	u8_pt q = rsa_default_malloc();
	u8_pt r = rsa_default_malloc();
	bi_div(a, b, r, q);
	
	i32_t x1 = 0, y1 = 0;
	rsa_extended_euclid_gcd(b, q, d, l, k, &x1, &y1);
	bi_mul(k, r, r);
	*markX = y1;
	if(x1) // negative
	{
		if(y1 == 1 && bi_comp(r, l) >= 0) 
		{
			bi_sub(r, l, l);
			*markY = 0;
		}
		else
		{
			if(y1 == 1) bi_sub(l, r, l);
			else bi_add(l, r, l);
			*markY = 1;
		}
	}
	else // positive
	{
		if(y1 == 0 && bi_comp(l, r) < 0) 
		{
			bi_sub(r, l, l);
			*markY = 1;
		}
		else
		{
			if(y1 == 0) bi_sub(l, r, l);
			else bi_add(l, r, l);
			*markY = 0;
		}
	}
	rsa_free(q);
	rsa_free(r);
}

i32_t rsa_calculate_keys(i8_t bits, u8_pt a, u8_pt b, u8_pt eOut, u8_pt dOut)
{
	u8_pt n		= rsa_default_malloc();
	u8_pt phi_n = rsa_default_malloc();
	u8_pt e		= rsa_default_malloc();
	u8_pt d		= rsa_default_malloc();
	u8_pt i		= rsa_default_malloc();
	u8_pt j		= rsa_default_malloc();

	bi_set_bit(phi_n, 0, 1);
	bi_set_bit(e, 0, 1);
	bi_mul(a, b, n);
	bi_sub(a, phi_n, phi_n);
	bi_sub(b, e, e);
	bi_mul(phi_n, e, phi_n);

	while(1)
	{
		i32_t tbits = (bits >> 1);
		rsa_rand_prime(rsa_max(8, tbits), e);
		if(rsa_relatively_prime(e, n) && rsa_relatively_prime(e, phi_n))
		{
			break;
		}
	}

	i32_t x = 0, y = 0;
	rsa_extended_euclid_gcd(e, phi_n, d, i, j, &x, &y);
	if(bi_bits(d) != 1)
	{
		// ERROR HERE, IT SHOULD BE ONE
		exit(-1);
	}
	if(x == 1)
	{
		rsa_free(n);
		rsa_free(phi_n);
		rsa_free(e);
		rsa_free(d);
		rsa_free(i);
		rsa_free(j);
		return 1;
	}
	bi_mul(e, i, j);
	bi_mod(j, phi_n, j);
	if(bi_bits(j) != 1)
	{
		// ERROR HERE, IT SHOULD BE ONE
		exit(-1);
	}
	bi_copy_bytes(eOut, e, 0, MAX_BYTES);
	bi_copy_bytes(dOut, i, 0, MAX_BYTES);

	rsa_free(n);
	rsa_free(phi_n);
	rsa_free(e);
	rsa_free(d);
	rsa_free(i);
	rsa_free(j);
	return 0;
}

void rsa_gen_keys(i32_t bits, rsa_key *keys)
{

	while(1)
	{
		rsa_gen_p_q(bits, keys->p, keys->q);
		i32_t retv = rsa_calculate_keys(bits, keys->p, keys->q, keys->e, keys->d);
		if(retv == 0)
		{
			break;
		}
	}
	bi_mul(keys->p, keys->q, keys->n);
}

void rsa_encrypt_bigint(u8_pt msg, u8_pt e, u8_pt n, u8_pt out)
{
	if(bi_bits(msg) > bi_bits(n))
	{
		// Error handling here
		bi_reset(out);
		return;
	}
	bi_powm(msg, e, n, out);
}


void rsa_decrypt_bigint(u8_pt msg, u8_pt d, u8_pt n, u8_pt out)
{
	rsa_encrypt_bigint(msg, d, n, out);
}

void rsa_sign_bigint(u8_pt msg, u8_pt d, u8_pt n, u8_pt out)
{
	rsa_decrypt_bigint(msg, d, n, out);
}

void rsa_verify_bigint(u8_pt msg, u8_pt e, u8_pt n, u8_pt out)
{
	rsa_encrypt_bigint(msg, e, n, out);
}

i32_t rsa_chopstring(u8_pt msg, u8_pt key, u8_pt n, u8_pt out, i32_t len, funcref func)
{
	i32_t nbits   = bi_bits(n) - 1;
	i32_t sbytes  = (nbits + 7) / 8,                nbytes = (nbits >> 3);
	i32_t nblocks = (len + nbytes - 1) / nbytes,    rbytes = len % nbytes, i = 0;

	u8_pt t = rsa_default_malloc();
	for(i=0; i<nblocks; i++)
	{
		i32_t offset = i * nbytes;
		if(i == nblocks - 1 && rbytes != 0)
			bi_copy_bytes(t, msg + offset, 0, rbytes);
		else
			bi_copy_bytes(t, msg + offset, 0, nbytes);
		func(t, key, n, t);
		memcpy(out + i * sbytes, t, sbytes);
	}
	rsa_free(t);
	return i * sbytes;
}

i32_t rsa_gluechops(u8_pt chops, u8_pt e, u8_pt n, u8_pt out, i32_t len, funcref func)
{
	i32_t nbits   = bi_bits(n) - 1;
	i32_t sbytes  = (nbits + 7) / 8, nbytes = nbits / 8;
	i32_t nblocks = len/sbytes, i = 0, wbits = 0, wbytes = 0;

	u8_pt t = rsa_default_malloc();
	for(i=0; i<nblocks; i++)
	{
		bi_copy_bytes(t, chops + i * sbytes, 0, sbytes);
		func(t, e, n, t);
		if(i == nblocks - 1)
		{
			wbits  = bi_bits(t);
			wbytes = (wbits + 7)/8;
			memcpy(out + i * nbytes, t, wbytes);
		}
		else
			memcpy(out + i * nbytes, t, nbytes);
	}
	rsa_free(t);
	return (i-1) * nbytes + wbytes;
}

